#ifndef CHAOSSOUP_HASH_H
#define CHAOSSOUP_HASH_H

#include <vector>
#include <array>
#include <numeric>
#include "gf.h"

// Hello from ChaosSoup-1!
// 6aa27c833b4ba0279b96029c9b8482e34801dc721b2fa24ebbcee523fdb778b3
// here is a hash-256 algorithm
namespace ChaosSoup {
    using u8  = std::uint8_t;
    using u16 = std::uint16_t;
    using u32 = std::uint32_t;
    using sz = std::size_t;
    using data_stream = std::vector<u8>;
    constexpr u8 block_size = 32;
    constexpr u8 block_size_mask = block_size - 1;
    using block_t = std::array<u8, block_size>;

    // a non-linear table
    constexpr std::array<u8, 256> sbox = {
        0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5, 0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
        0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0, 0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0,
        0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC, 0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15,
        0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A, 0x07, 0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75,
        0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0, 0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84,
        0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B, 0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF,
        0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85, 0x45, 0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8,
        0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5, 0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2,
        0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17, 0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73,
        0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88, 0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB,
        0xE0, 0x32, 0x3A, 0x0A, 0x49, 0x06, 0x24, 0x5C, 0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79,
        0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9, 0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08,
        0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6, 0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A,
        0x70, 0x3E, 0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E, 0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E,
        0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94, 0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF,
        0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68, 0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB, 0x16
    };

    // 非线性索引表（预计算）
    /*
     * idx_table:
        0f 13 19 14 0c 06 02 1d 16 01 0b 0e 01 0d 1e 18
        13 12 10 12 05 12 0f 0e 18 0e 16 1a 0f 10 1c 18
        01 02 01 02 10 18 09 15 12 11 19 0f 1f 03 17 17
        04 1f 07 1b 1b 04 05 09 07 10 10 1e 16 03 04 1b
        08 13 09 19 18 03 11 14 18 1c 0c 05 0c 1f 0a 14
        19 0b 00 10 04 03 07 10 07 12 09 1e 03 05 13 16
        0a 12 1f 04 0b 04 15 15 0d 06 02 10 1a 1b 0c 1d
        1b 17 08 1e 00 0e 1f 0b 0b 00 01 05 12 00 0d 08
        14 0d 11 11 14 05 0c 15 1c 13 11 1a 08 16 1a 1d
        0c 11 06 07 06 0f 02 19 0e 13 0f 16 05 15 0a 00
        1c 14 1d 0b 00 06 00 17 1a 09 19 0e 03 07 18 16
        1b 11 11 00 1c 0f 07 1c 01 1c 0a 17 09 15 1b 09
        0d 17 01 0b 1f 12 02 1e 15 06 1a 1c 02 0a 1a 1b
        1e 19 03 0a 01 03 08 0f 0d 13 1d 0e 16 19 1e 0d
        1d 07 0b 13 04 02 1f 06 08 1d 17 14 17 1f 0d 04
        1d 15 18 0c 08 1a 0a 05 09 0a 08 0e 06 1e 0c 14
     * each one appeared 8 times
     */
    constexpr std::array<u8, 256> idx_table = []() consteval noexcept {
        std::array<u8, 256> table{};
        for (u16 i = 0; i < 256; ++i) {
            // 增强非线性的索引计算
            const u8 v = sbox[i];
            table[i] = (v & block_size_mask) ^ ((v>>3) & block_size_mask);
        }
        return table;
    }();

    // 非线性索引函数
    constexpr inline u8 get_index(const sz i) noexcept {
        return idx_table[i & 0xFF];
    }

    constexpr block_t rs = {
        5, 4, 1, 2, 7, 4, 2, 4,
        5, 4, 4, 2, 3, 4, 4, 3,
        2, 3, 7, 7, 1, 3, 5, 6,
        2, 4, 7, 2, 3, 1, 4, 7
    };

    // Add-Rotate-XOR
    inline constexpr u8 ARX(const u8 x, const u8 y, const u8 r1, const u8 r2, const u8 u) noexcept {
        return sbox[static_cast<u8>(
            (x + std::rotl(y, r1)) ^ (y + std::rotr(x, r2))
        )] + u;
    }

    // Mix With Error
    constexpr inline u8 MWE(const u8 x, const u8 y, const u8 e) noexcept {
        return gf_mul(
            sbox[x ^ std::rotr(y, e&7)], (y + sbox[e])
        ) ^ std::rotl(x,e>>5) ^ (e & 0x18);
    }

    constexpr std::array<u32, 8> ChaChaConstants {
        // 原始 ChaCha 常量（保持不动）
        0x61707865u,  // "expa" (小端：0x65 0x78 0x70 0x61 -> "a p x e")
        0x3320646eu,  // "nd 3" (小端：0x6e 0x64 0x20 0x33 -> "3   n d")
        0x79622d32u,  // "2-by" (小端：0x32 0x2d 0x62 0x79 -> "y - b 2")
        0x6b206574u,  // "te k" (小端：0x74 0x65 0x20 0x6b -> "k   e t")

        // 新增部分（严格小端编码）
        0x65792077u,  // "ey w" (小端：0x77 0x20 0x79 0x65 -> "e y   w")
        0x69746820u,  // "ith " (小端：0x20 0x68 0x74 0x69 -> "i t h  ")
        0x6f616843u,  // "Chao" (小端：0x43 0x68 0x61 0x6f -> "o a h C")
        0x70756f53u   // "Soup" (小端：0x53 0x6f 0x75 0x70 -> "p u o S")
    };

    // quarter round from ChaCha
    constexpr inline void QuarterRound(
        u32& a, u32& b,
        u32& c, u32& d
    ) noexcept {
        a += b; d ^= a; d = std::rotl(d, 16);
        c += d; b ^= c; b = std::rotl(b, 12);
        a += b; d ^= a; d = std::rotl(d, 8);
        c += d; b ^= c; b = std::rotl(b, 7);
    }

    constexpr block_t ChaCha_transform(const block_t& input) noexcept {
        // 将输入转换为状态字
        std::array<u32, 8> state = [&input]() constexpr noexcept {
            std::array<u32, 8> s{};
            for (sz i = 0; i < s.size(); ++i) {
                s[i] = std::accumulate(
                    input.begin() + i * 4,
                    input.begin() + (i + 1) * 4, 0u,
                    [](const u32 acc, const u8 b) -> u32 {
                        return (acc << 8) | b;
                    });
            }
            return s;
        }();

        // 添加常量
        for (u8 i = 0; i < 8; ++i) {
            state[i] ^= ChaChaConstants[i];
        }

        // 16 轮双向扩散
        // 0 1 2 3
        // 4 5 6 7
        for (u8 _ = 0; _ < 8; ++_) {
            // 列轮
            QuarterRound(state[0], state[1], state[2], state[3]);
            QuarterRound(state[4], state[5], state[6], state[7]);
            // 对角轮
            QuarterRound(state[0], state[5], state[2], state[7]);
            QuarterRound(state[4], state[1], state[6], state[3]);
        }

        // 转换回字节
        block_t output;
        for (u8 i = 0; i < 8; ++i) {
            output[i*4]   = (state[i] >> 24) & 0xFF;
            output[i*4+1] = (state[i] >> 16) & 0xFF;
            output[i*4+2] = (state[i] >> 8)  & 0xFF;
            output[i*4+3] =  state[i]        & 0xFF;
        }
        return output;
    }

    inline data_stream pad(const data_stream& data) noexcept {
        constexpr u8 pad_block_size = block_size * 2;
        const sz data_bytes = data.size();
        const sz bit_len = data_bytes * 8; // 原始消息的位长度

        // 计算填充后的总长度需满足 (data_bytes + 1 + pad_zero_bytes + 8) % 64 == 0
        const sz rem = (data_bytes + 1) % pad_block_size;
        const sz pad_zero_bytes = (56 - rem) % pad_block_size; // 计算需要填充的0字节数

        data_stream padded = data;


        // 1. 附加 0x80
        padded.push_back(0x80);

        // 2. 补 0x00
        padded.insert(padded.end(), pad_zero_bytes, 0x00);

        // 3. 附加位长度（64 位大端序）
        for (u8 i = 8; i-- > 0; ) {
            padded.push_back(static_cast<u8>((bit_len >> (i * 8)) & 0xFF));
        }

        return padded;
    }

    /*
         * 初始向量生成说明：
         * 基于诗句"泠泠风不停，点点星烛明。"(UTF-8编码)构建
         * 采用以下保护措施：
         * 1. 循环使用诗句完整UTF-8编码
         * 2. 应用旋转和非线性变换
         * 2. 非线性压缩
         * 整个过程可完全复现验证，确保无后门
         */
    constexpr block_t init_state() noexcept {
        // "泠泠风不停，点点星烛明。" 的UTF-8编码 (共30字节)
        constexpr u8 poem[] = {
            0xE6, 0xB3, 0xA0, // 泠
            0xE6, 0xB3, 0xA0, // 泠
            0xE9, 0xA3, 0x8E, // 风
            0xE4, 0xB8, 0x8D, // 不
            0xE5, 0x81, 0x9C, // 停
            0xEF, 0xBC, 0x8C, // ，
            0xE7, 0x82, 0xB9, // 点
            0xE7, 0x82, 0xB9, // 点
            0xE6, 0x98, 0x9F, // 星
            0xE7, 0x83, 0x9B, // 烛
            0xE6, 0x98, 0x8E, // 明
            0xEF, 0xBC, 0x82  // 。
        };
        block_t s{};

        for (u8 i = 0; i < block_size; ++i) {
            // 每3个输入字节产生1个输出字节
            const u8 idx1 = (i * 3) % sizeof(poem),
                  idx2 = (i * 3 + 1) % sizeof(poem),
                  idx3 = (i * 3 + 2) % sizeof(poem);
            s[i] = MWE(poem[idx1], poem[idx2], poem[idx3]);
        }
        return s;
    };

    // Hash function that produces a 32-byte (256-bit) hash
    // from ChaosSoup-1
    inline block_t hash(const data_stream& input, block_t state = init_state()) noexcept {
        data_stream C(pad(input));
        const sz size = C.size();

        // Process the input in multiple rounds for better diffusion
        for (u8 round = 0; round < 20; ++round) {
            for (sz i = 0; i < size; ++i) {
                const u8 i_b = get_index(i),
                         i4k_b = get_index(i-4),
                         i2k_b = get_index(i+2);
                const sz i4k_s = (i - 4) % size,
                         i2k_s = (i + 2) % size;

                // 非马尔可夫更新机制
                C[i] ^= C[i4k_s];
                state[i_b] = sbox[state[i_b]] ^
                             MWE(state[i2k_b], C[i], C[i4k_s] ^ state[i4k_b]);

                // if i % 16 == 0: transform state
                if ((i & 0xF) == 0) {
                    ChaCha_transform(state);
                }

                // ARX
                C[i] = ARX(C[i], state[i_b], rs[i_b], rs[i2k_b], state[i4k_b]);
                // XOR with future
                C[i] ^= C[i2k_s];
                // MWE
                C[i] = MWE(C[i], state[i_b], state[i2k_b]);
            }
        }

        // Final compression to 32 bytes
        block_t output(state);
        for (sz i = 0; i < size; ++i) {
            const u8 index = get_index(i);
            output[index] = sbox[output[index] ^ C[i]];
        }

        return ChaCha_transform(output);
    }
}

#endif //CHAOSSOUP_HASH_H